import numpy as np
from scipy.spatial import KDTree

class CornerPointGrid:
    def __init__(self, filename):
        self.nx = 0
        self.ny = 0
        self.nz = 0
        self.coords = []  # COORD数据
        self.zcorn = []   # ZCORN数据
        self.read_grid(filename)

    def read_grid(self, filename):
        """读取角点网格文件"""
        with open(filename, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        # 读取网格维度
        for i, line in enumerate(lines):
            if 'SPECGRID' in line:
                dims = lines[i+1].split()
                self.nx = int(dims[0])
                self.ny = int(dims[1])
                self.nz = int(dims[2])
                break

        # 读取坐标数据
        reading_coords = False
        reading_zcorn = False
        
        for line in lines:
            line = line.strip()
            if not line or line.startswith('/'):
                continue

            if line == 'COORD':
                reading_coords = True
                reading_zcorn = False
                continue
            elif line == 'ZCORN':
                reading_coords = False
                reading_zcorn = True
                continue

            if reading_coords:
                try:
                    values = [float(x) for x in line.split()]
                    self.coords.append(values)
                except ValueError:
                    reading_coords = False
                    
            if reading_zcorn:
                try:
                    values = [float(x) for x in line.split()]
                    self.zcorn.extend(values)
                except ValueError:
                    reading_zcorn = False

        self.coords = np.array(self.coords)
        self.zcorn = np.array(self.zcorn)

    def get_cell_vertices(self, i, j, k):
        """获取指定单元的8个顶点坐标"""
        # 计算在当前层的索引
        nx_lines = self.nx + 1
        ny_lines = self.ny + 1
        
        # 获取COORD中的坐标点
        pillar_indices = [
            j * nx_lines + i,
            j * nx_lines + (i + 1),
            (j + 1) * nx_lines + i,
            (j + 1) * nx_lines + (i + 1)
        ]
        
        # 获取每个柱子的顶部和底部坐标
        pillars = []
        for idx in pillar_indices:
            pillars.append(self.coords[idx])
            
        # 从ZCORN获取z坐标
        z_indices = []
        nxy = self.nx * self.ny
        layer_offset = k * nxy * 8
        
        # 计算8个顶点的z值索引
        for corner in range(8):
            z_idx = layer_offset + (j * self.nx + i) * 2 + (corner // 4) * nxy * 2 + (corner % 4)
            z_indices.append(z_idx)
            
        # 构建顶点坐标
        vertices = []
        for p_idx, pillar in enumerate(pillars):
            # 底部顶点
            bottom_z = self.zcorn[z_indices[p_idx]]
            # 顶部顶点
            top_z = self.zcorn[z_indices[p_idx + 4]]
            
            # 计算实际的xyz坐标
            x = pillar[0]
            y = pillar[1]
            vertices.append([x, y, bottom_z])
            vertices.append([x, y, top_z])
            
        return np.array(vertices)

    def calculate_cell_centers(self):
        """计算所有网格单元的中心点坐标"""
        centers = []
        
        for k in range(self.nz):
            for j in range(self.ny):
                for i in range(self.nx):
                    vertices = self.get_cell_vertices(i, j, k)
                    center = np.mean(vertices, axis=0)
                    centers.append(center)
                    
        return np.array(centers)

def read_props(filename):
    """读取属性文件，支持多个属性"""
    props = {}
    current_prop = None
    current_values = []
    
    with open(filename, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('/'):
                continue
                
            # 检查是否是属性名
            if line in ['YOUNGMOD', 'POISSONR']:
                # 如果已经在读取某个属性，先保存它
                if current_prop:
                    props[current_prop] = np.array(current_values)
                    current_values = []
                current_prop = line
                continue
            
            # 处理属性值行
            try:
                # 处理形如 "1*123.456" 的数据
                values = []
                for item in line.split():
                    if '*' in item:
                        count, value = item.split('*')
                        values.extend([float(value)] * int(count))
                    else:
                        values.append(float(item))
                current_values.extend(values)
            except ValueError:
                continue
    
    # 保存最后一个属性
    if current_prop:
        props[current_prop] = np.array(current_values)
    
    return props

def interpolate_properties(centers_a, centers_b, props_a):
    """对每个属性进行插值"""
    props_b = {}
    
    # 构建KD树
    tree = KDTree(centers_a)
    
    # 对每个属性进行插值
    for prop_name, prop_values in props_a.items():
        # 找到最近的8个点
        distances, indices = tree.query(centers_b, k=8)
        
        # 计算权重（使用距离的倒数）
        weights = 1.0 / (distances + 1e-10)
        weights = weights / weights.sum(axis=1, keepdims=True)
        
        # 插值计算
        prop_b = np.sum(prop_values[indices] * weights, axis=1)
        props_b[prop_name] = prop_b
    
    return props_b

def write_props(filename, props):
    """写入属性文件，支持多个属性"""
    with open(filename, 'w') as f:
        for prop_name, values in props.items():
            f.write(prop_name + '\n')
            # 每行写入6个值，保持格式一致
            for i in range(0, len(values), 6):
                line_values = values[i:i+6]
                line = ' '.join([f'1*{v:.8f}' for v in line_values])
                f.write(' ' + line + '\n')
            f.write('/\n')

def main():
    # 1. 读取网格
    grid_a = CornerPointGrid("Grid_A.out")
    grid_b = CornerPointGrid("Grid_B.out")
    
    # 2. 读取属性
    props_a = read_props("A_props.txt")
    
    # 3. 计算网格中心点
    centers_a = grid_a.calculate_cell_centers()
    centers_b = grid_b.calculate_cell_centers()
    
    # 4. 插值计算B网格的属性
    props_b = interpolate_properties(centers_a, centers_b, props_a)
    
    # 5. 输出结果
    write_props("B_props.txt", props_b)

if __name__ == "__main__":
    main()